import prior

OBJNAV_TYPES_THOR = [
    "AlarmClock",
    "Apple",
    # "BaseballBat",
    "BasketBall",
    "Bed",
    "Bowl",
    "Chair",
    "GarbageCan",
    "HousePlant",
    "Laptop",
    "Mug",
    "Sofa",
    "SprayBottle",
    "Television",
    "Toilet",
    "Vase",
]

PICKUP_TYPES_THOR = [
    "AlarmClock",
    "Apple",
    # "BaseballBat",
    "BasketBall",
    "Bowl",
    "Laptop",
    "Mug",
    "SprayBottle",
    "Vase",
]

RECEPTACLE_TYPES_THOR = [
    "ArmChair",
    "Bed",
    # "Bowl",
    # "Cabinet",
    "Chair",
    "CoffeeTable",
    "CounterTop",
    "Desk",
    "DiningTable",
    # "Drawer",
    "Dresser",
    # "Fridge",
    # "GarbageCan",
    # "Microwave",
    # "Pan",
    # "Plate",
    # "Pot",
    "Shelf",
    # "ShelvingUnit",
    "SideTable",
    "Sofa",
    "Stool",
    "TVStand",
]

ALL_PICKUP_TYPES_THOR = [
    "AlarmClock",
    "AluminumFoil",
    "Apple",
    "AppleSliced",
    # "BaseballBat",
    "Book",
    "Boots",
    "Bottle",
    "Bowl",
    "Box",
    "Bread",
    "BreadSliced",
    "ButterKnife",
    "Candle",
    "CD",
    "CellPhone",
    "Cloth",
    "CreditCard",
    "Cup",
    "DishSponge",
    "Dumbbell",
    "Egg",
    "EggCracked",
    "Fork",
    "HandTowel",
    "Kettle",
    "KeyChain",
    "Knife",
    "Ladle",
    "Laptop",
    "Lettuce",
    "LettuceSliced",
    "Mug",
    "Newspaper",
    "Pan",
    "PaperTowelRoll",
    "Pen",
    "Pencil",
    "PepperShaker",
    "Pillow",
    "Plate",
    "Plunger",
    "Pot",
    "Potato",
    "PotatoSliced",
    "RemoteControl",
    "SaltShaker",
    # "Sandwich", #rose: i think this should be pickupable
    "ScrubBrush",
    "SoapBar",
    "SoapBottle",
    "Spatula",
    "Spoon",
    "SprayBottle",
    "Statue",
    "TableTopDecor",
    "TeddyBear",
    "TennisRacket",
    "TissueBox",
    "ToiletPaper",
    "Tomato",
    "TomatoSliced",
    "Towel",
    "Vase",
    "Watch",
    "WateringCan",
    "WineBottle",
]

AI2THOR_OBJECT_TYPE_TO_WORDNET_SYNSET = {
    "AlarmClock": "alarm_clock.n.01",
    "AluminumFoil": "aluminum_foil.n.01",
    "Apple": "apple.n.01",
    "AppleSliced": "apple.n.01",
    "ArmChair": "armchair.n.01",
    "BaseballBat": "baseball_bat.n.01",
    "BasketBall": "basketball.n.02",
    "Bathtub": "bathtub.n.01",
    "BathtubBasin": "bathtub.n.01",
    "Bed": "bed.n.01",
    "Blinds": "blind.n.03",
    "Book": "book.n.02",
    "Boots": "boot.n.01",
    "Bottle": "bottle.n.01",
    "Bowl": "bowl.n.03",
    "Box": "carton.n.02",
    "Bread": "bread.n.01",
    "BreadSliced": "bread.n.01",
    "ButterKnife": "butter_knife.n.01",
    "CD": "compact_disk.n.01",
    "Cabinet": "cabinet.n.01",
    "Candle": "candle.n.01",
    "Cart": "handcart.n.01",
    "CellPhone": "cellular_telephone.n.01",
    "Chair": "straight_chair.n.01",
    "Cloth": "fabric.n.01",
    "ClothesDryer": "clothes_dryer.n.01",
    "CoffeeMachine": "coffee_maker.n.01",
    "CoffeeTable": "coffee_table.n.01",
    "CounterTop": "countertop.n.01",
    "CreditCard": "credit_card.n.01",
    "Cup": "cup.n.01",
    "Curtains": "curtain.n.01",
    "Desk": "desk.n.01",
    "DeskLamp": "table_lamp.n.01",
    "Desktop": "desktop_computer.n.01",
    "DiningTable": "dining_table.n.01",
    "DishSponge": "sponge.n.01",
    "DogBed": "pad.n.04",
    "Doorframe": "doorframe.n.01",
    "Doorway": "doorway.n.01",
    "Drawer": "drawer.n.01",
    "Dresser": "chest_of_drawers.n.01",
    "Dumbbell": "dumbbell.n.01",
    "Egg": "egg.n.02",
    "EggCracked": "egg.n.02",
    "Faucet": "faucet.n.01",
    "Floor": "floor.n.01",
    "FloorLamp": "floor_lamp.n.01",
    "Footstool": "footstool.n.01",
    "Fork": "fork.n.01",
    "Fridge": "refrigerator.n.01",
    "GarbageBag": "bin_liner.n.01",
    "GarbageCan": "ashcan.n.01",
    "HandTowel": "hand_towel.n.01",
    "HandTowelHolder": "towel_rack.n.01",
    "HousePlant": "houseplant.n.01",
    "Kettle": "kettle.n.01",
    "KeyChain": "key_ring.n.01",
    "Knife": "knife.n.01",
    "Ladle": "ladle.n.01",
    "Laptop": "laptop.n.01",
    "LaundryHamper": "clothes_hamper.n.01",
    "Lettuce": "lettuce.n.03",
    "LettuceSliced": "lettuce.n.03",
    "LightSwitch": "switch.n.01",
    "Microwave": "microwave.n.02",
    "Mirror": "mirror.n.01",
    "Mug": "mug.n.04",
    "Newspaper": "newspaper.n.03",
    "Ottoman": "footstool.n.01",
    "Painting": "painting.n.01",
    "Pan": "pan.n.01",
    "PaperTowelRoll": "paper_towel.n.01",
    "Pen": "pen.n.01",
    "Pencil": "pencil.n.01",
    "PepperShaker": "pepper_shaker.n.01",
    "Pillow": "pillow.n.01",
    "Plate": "plate.n.04",
    "Plunger": "plunger.n.03",
    "Poster": "poster.n.01",
    "Pot": "pot.n.01",
    "Potato": "potato.n.01",
    "PotatoSliced": "potato.n.01",
    "RemoteControl": "remote_control.n.01",
    "RoomDecor": "decoration.n.01",
    "Safe": "safe.n.01",
    "SaltShaker": "saltshaker.n.01",
    "ScrubBrush": "scrub_brush.n.01",
    "Shelf": "shelf.n.01",
    "ShelvingUnit": "shelf.n.01",
    "ShowerCurtain": "shower_curtain.n.01",
    "ShowerDoor": "door.n.01",
    "ShowerGlass": "door.n.01",
    "ShowerHead": "showerhead.n.01",
    "SideTable": "stand.n.04",
    "Sink": "sink.n.01",
    "SinkBasin": "sink.n.01",
    "SoapBar": "bar_soap.n.01",
    "SoapBottle": "soap_dispenser.n.01",
    "Sofa": "sofa.n.01",
    "Spatula": "spatula.n.01",
    "Spoon": "spoon.n.01",
    "SprayBottle": "atomizer.n.01",
    "Statue": "statue.n.01",
    "Stool": "stool.n.01",
    "StoveBurner": "burner.n.02",
    "StoveKnob": "knob.n.02",
    "TVStand": "stand.n.04",
    "TableTopDecor": "knickknack.n.01",
    "TeddyBear": "teddy.n.01",
    "Television": "television_receiver.n.01",
    "TennisRacket": "tennis_racket.n.01",
    "TissueBox": "tissue.n.02",
    "Toaster": "toaster.n.02",
    "Toilet": "toilet.n.02",
    "ToiletPaper": "toilet_tissue.n.01",
    "ToiletPaperHanger": "hanger.n.02",
    "Tomato": "tomato.n.01",
    "TomatoSliced": "tomato.n.01",
    "Towel": "towel.n.01",
    "TowelHolder": "towel_rack.n.01",
    "VacuumCleaner": "vacuum.n.04",
    "Vase": "vase.n.01",
    "Wall": "wall.n.01",
    "WashingMachine": "washer.n.03",
    "Watch": "watch.n.01",
    "WateringCan": "watering_can.n.01",
    "Window": "window.n.01",
    "WineBottle": "wine_bottle.n.01",
}

AI2THOR_OBJECT_TYPE_TO_MOST_SPECIFIC_WORDNET_LEMMA = {
    "AlarmClock": "alarm_clock",
    "AluminumFoil": "aluminum_foil",
    "Apple": "apple",
    "AppleSliced": "apple",
    "ArmChair": "armchair",
    "BaseballBat": "baseball_bat",
    "BasketBall": "basketball",
    "Bathtub": "bathtub",
    "BathtubBasin": "bathtub",
    "Bed": "bed",
    "Blinds": "blind",
    "Book": "book",
    "Boots": "boot",
    "Bottle": "bottle",
    "Bowl": "bowl",
    "Box": "box",
    "Bread": "bread",
    "BreadSliced": "bread",
    "ButterKnife": "butter_knife",
    "CD": "compact_disk",
    "Cabinet": "cabinet",
    "Candle": "candle",
    "Cart": "handcart",
    "CellPhone": "cellular_telephone",
    "Chair": "chair",
    "Cloth": "fabric",
    "ClothesDryer": "clothes_dryer",
    "CoffeeMachine": "coffee_maker",
    "CoffeeTable": "coffee_table",
    "CounterTop": "countertop",
    "CreditCard": "credit_card",
    "Cup": "cup",
    "Curtains": "curtain",
    "Desk": "desk",
    "DeskLamp": "table_lamp",
    "Desktop": "desktop_computer",
    "DiningTable": "dining_table",
    "DishSponge": "sponge",
    "DogBed": "pad",
    "Doorframe": "doorframe",
    "Doorway": "doorway",
    "Drawer": "drawer",
    "Dresser": "chest_of_drawers",
    "Dumbbell": "dumbbell",
    "Egg": "egg",
    "EggCracked": "egg",
    "Faucet": "faucet",
    "Floor": "flooring",
    "FloorLamp": "floor_lamp",
    "Footstool": "footstool",
    "Fork": "fork",
    "Fridge": "refrigerator",
    "GarbageBag": "bin_liner",
    "GarbageCan": "ashcan",
    "HandTowel": "hand_towel",
    "HandTowelHolder": "towel_rack",
    "HousePlant": "houseplant",
    "Kettle": "boiler",
    "KeyChain": "key_ring",
    "Knife": "knife",
    "Ladle": "ladle",
    "Laptop": "laptop",
    "LaundryHamper": "clothes_hamper",
    "Lettuce": "lettuce",
    "LettuceSliced": "lettuce",
    "LightSwitch": "electric_switch",
    "Microwave": "microwave_oven",
    "Mirror": "mirror",
    "Mug": "mug",
    "Newspaper": "newspaper",
    "Ottoman": "pouffe",
    "Painting": "painting",
    "Pan": "cooking_pan",
    "PaperTowelRoll": "paper_towel",
    "Pen": "pen",
    "Pencil": "pencil",
    "PepperShaker": "pepper_shaker",
    "Pillow": "pillow",
    "Plate": "plate",
    "Plunger": "plumber's_helper",
    "Poster": "placard",
    "Pot": "pot",
    "Potato": "Irish_potato",
    "PotatoSliced": "Irish_potato",
    "RemoteControl": "remote_control",
    "RoomDecor": "decoration",
    "Safe": "safe",
    "SaltShaker": "saltshaker",
    "ScrubBrush": "scrub_brush",
    "Shelf": "shelf",
    "ShelvingUnit": "shelf",
    "ShowerCurtain": "shower_curtain",
    "ShowerDoor": "door",
    "ShowerGlass": "door",
    "ShowerHead": "showerhead",
    "SideTable": "stand",
    "Sink": "sink",
    "SinkBasin": "sink",
    "SoapBar": "bar_soap",
    "SoapBottle": "soap_dispenser",
    "Sofa": "sofa",
    "Spatula": "spatula",
    "Spoon": "spoon",
    "SprayBottle": "atomizer",
    "Statue": "statue",
    "Stool": "stool",
    "StoveBurner": "burner",
    "StoveKnob": "knob",
    "TVStand": "stand",
    "TableTopDecor": "knickknack",
    "TeddyBear": "teddy_bear",
    "Television": "television",
    "TennisRacket": "tennis_racket",
    "TissueBox": "tissue_paper",
    "Toaster": "toaster",
    "Toilet": "crapper",
    "ToiletPaper": "toilet_tissue",
    "ToiletPaperHanger": "hanger",
    "Tomato": "tomato",
    "TomatoSliced": "tomato",
    "Towel": "towel",
    "TowelHolder": "towel_rack",
    "VacuumCleaner": "vacuum_cleaner",
    "Vase": "vase",
    "Wall": "wall",
    "WashingMachine": "automatic_washer",
    "Watch": "watch",
    "WateringCan": "watering_can",
    "Window": "window",
    "WineBottle": "wine_bottle",
}

_cached_bad_asset_ids = None


def bad_asset_ids():
    global _cached_bad_asset_ids
    if _cached_bad_asset_ids is None:
        _cached_bad_asset_ids = prior.load_dataset(
            dataset="spoc-data", entity="spoc-robot", revision="spoc-additional-annotation"
        )
    return _cached_bad_asset_ids
